package kelvin.pagerank;

import java.io.IOException;
import java.util.StringTokenizer;
import java.util.ArrayList;
import java.util.Collections;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.IntWritable;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.FloatWritable;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.mapreduce.Job;
import org.apache.hadoop.mapreduce.Mapper;
import org.apache.hadoop.mapreduce.Reducer;
import org.apache.hadoop.mapreduce.Partitioner;
import org.apache.hadoop.mapreduce.lib.input.FileInputFormat;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;
import org.apache.hadoop.mapreduce.lib.input.TextInputFormat;
import org.apache.hadoop.mapreduce.lib.output.TextOutputFormat;
import org.apache.hadoop.util.GenericOptionsParser;

public class PageRank {

  public static class StartupMapper extends Mapper<LongWritable, Text, Text, Text> {

    private float initPageRank;

    public void setup(Context context) {
      initPageRank = context.getConfiguration().getFloat("initPageRank", 1.0f);
    }

    private int tokenAmount(String str) {
      int result = 0;
      StringTokenizer itr = new StringTokenizer(str.toString());
      while (itr.hasMoreTokens()) {
        itr.nextToken();
        result++;
      }
      return result;
    }

    /* 
     * input: (line-offset, line)
     * output: (source-node, =init-pagerank #degree adj-list)
     */
    public void map(LongWritable lineOffset, Text line, Context context) throws IOException, InterruptedException {
      int degree = tokenAmount(line.toString()) - 1;
      StringTokenizer itr = new StringTokenizer(line.toString());
      if (!itr.hasMoreTokens()) return;
      String key = itr.nextToken();
      String value = String.format("=%f #%d", initPageRank, degree);
      while (itr.hasMoreTokens())
        value += " " + itr.nextToken();
      context.write(new Text(key), new Text(value));
    }
  }

  public static class IterationMapper extends Mapper<LongWritable, Text, Text, Text> {

    private float dampingFactor;

    public void setup(Context context) {
      dampingFactor = context.getConfiguration().getFloat("dampingFactor", 1.0f);
    }

    /* 
     * input: (line-offset, source-node =current-rank #degree adj-list)
     * output: (dest-node, +rank-given) ...,
     *         (source-node, #degree adj-list)
     */
    public void map(LongWritable lineOffset, Text line, Context context) throws IOException, InterruptedException {
      StringTokenizer itr = new StringTokenizer(line.toString());
      String src = itr.nextToken();
      String adjList = "";
      float rank = 0.0f;
      int degree = 1;
      while (itr.hasMoreTokens()) {
        String token = itr.nextToken();
        if (token.charAt(0) == '=') {
          rank = Float.parseFloat(token.substring(1));
        } else if (token.charAt(0) == '#') {
          degree = Integer.parseInt(token.substring(1));
          adjList += String.format("#%d", degree);
        } else {
          String dst = token;
          context.write(new Text(dst), new Text(String.format("+%f", dampingFactor * rank / degree)));
          adjList += " " + dst;
        }
      }
      context.write(new Text(src), new Text(String.format("+%f", (1.0f - dampingFactor) * rank)));
      context.write(new Text(src), new Text(adjList));
    }
  }
  
  public static class IterationReducer extends Reducer<Text, Text ,Text, Text> {
    /* 
     * input: (source-node, +rank-received)
     *        (source-node, #degree adj-list)
     * output: (source-node, =new-rank #degree adj-list)
     */
    public void reduce(Text key, Iterable<Text> values, Context context) throws IOException, InterruptedException {
      float rank = 0.0f;
      String adjList = "";
      for (Text v : values) {
        String value = v.toString();
        if (value.charAt(0) == '+')
          rank += Float.parseFloat(value.substring(1));
        else
          adjList = value;
      }
      context.write(key, new Text(String.format("=%f %s", rank, adjList)));
    }
  }

  public static class SortMapper extends Mapper<LongWritable, Text, FloatWritable, Text> {

    /* 
     * input: (line-offset, source-node =current-rank #degree adj-list)
     * output: (current-rank, source-node)
     */
    public void map(LongWritable lineOffset, Text line, Context context) throws IOException, InterruptedException {
      StringTokenizer itr = new StringTokenizer(line.toString());
      String source = itr.nextToken();
      float rank = Float.parseFloat(itr.nextToken().substring(1));
      context.write(new FloatWritable(rank), new Text(source));
    }
  }

  public static class SortPartitioner extends Partitioner<FloatWritable, Text> {
    public int getPartition(FloatWritable rank, Text node, int numPartition) {
      return Math.min(Math.max(0, (int)Math.floor(rank.get() / 10.0 * numPartition)), numPartition);
    }
  }

  public static class SortReducer extends Reducer<FloatWritable, Text, FloatWritable, Text> {

    /* 
     * input: (rank, node)
     * output: (rank, node) with nodes sorted
     */
    public void reduce(FloatWritable rank, Iterable<Text> nodes, Context context) throws IOException, InterruptedException {
      ArrayList<String> a = new ArrayList<String>();
      for (Text node : nodes)
        a.add(node.toString());
      Collections.sort(a);
      for (String node : a)
        context.write(rank, new Text(node));
    }
  }
  
  public static void main(String[] args) throws Exception {
    Configuration conf = new Configuration();
    args = new GenericOptionsParser(conf, args).getRemainingArgs();
    if (args.length != 5) {
      System.err.println("Usage: PageRank <init_pagerank> <damping_factor> <max_iteration_times> <input> <output>");
      System.exit(2);
    }

    float initPageRank = Float.parseFloat(args[0]);
    float dampingFactor = Float.parseFloat(args[1]);
    int maxIterationTimes = Integer.valueOf(args[2]);
    String inputDirName = args[3];
    String outputDirName = args[4];

    conf.setFloat("initPageRank", initPageRank);
    conf.setFloat("dampingFactor", dampingFactor);
    System.out.printf("initPageRank = %f\n", initPageRank);
    System.out.printf("dampingFactor = %f\n", dampingFactor);

    Job job;
    job = new Job(conf, "PageRank Startup");
    job.setJarByClass(PageRank.class);
    job.setMapperClass(StartupMapper.class);
    job.setOutputKeyClass(Text.class);
    job.setOutputValueClass(Text.class);
    job.setInputFormatClass(TextInputFormat.class);
    job.setOutputFormatClass(TextOutputFormat.class);
    FileInputFormat.addInputPath(job, new Path(inputDirName));
    FileOutputFormat.setOutputPath(job, new Path(outputDirName + "/iteration-0000"));
    System.out.printf("%s -> %s\n", inputDirName, outputDirName + "/iteration-0000");
    job.waitForCompletion(true);

    String lastOutDir = "";
    for (int i = 0; i < maxIterationTimes; i++) {
      String in = outputDirName + String.format("/iteration-%04d", i);
      String out = outputDirName + String.format("/iteration-%04d", i + 1);
      lastOutDir = out;
      System.out.printf("%s -> %s\n", in, out);
      job = new Job(conf, String.format("PageRank Iteration-%04d", i + 1));
      job.setJarByClass(PageRank.class);
      job.setMapperClass(IterationMapper.class);
      job.setReducerClass(IterationReducer.class);
      job.setOutputKeyClass(Text.class);
      job.setOutputValueClass(Text.class);
      job.setInputFormatClass(TextInputFormat.class);
      job.setOutputFormatClass(TextOutputFormat.class);
      FileInputFormat.addInputPath(job, new Path(in));
      FileOutputFormat.setOutputPath(job, new Path(out));
      job.waitForCompletion(true);
    }

    job = new Job(conf, "PageRank Sorting");
    job.setJarByClass(PageRank.class);
    job.setMapperClass(SortMapper.class);
    job.setPartitionerClass(SortPartitioner.class);
    job.setReducerClass(SortReducer.class);
    job.setOutputKeyClass(FloatWritable.class);
    job.setOutputValueClass(Text.class);
    job.setInputFormatClass(TextInputFormat.class);
    job.setOutputFormatClass(TextOutputFormat.class);
    FileInputFormat.addInputPath(job, new Path(lastOutDir));
    FileOutputFormat.setOutputPath(job, new Path(outputDirName + "/final"));
    System.out.printf("%s -> %s\n", lastOutDir, outputDirName + "/final");
    System.exit(job.waitForCompletion(true) ? 0 : 1);
  }
}
